import asyncio

from langchain.chat_models import init_chat_model
from langchain_core.output_parsers import JsonOutputParser
import os

key = os.getenv("OPENAI_API_KEY")
# print(key)
api_key = str(key)

llm = init_chat_model(
    model="gpt-4o-mini",
    model_provider="openai",
    base_url="https://api.zetatechs.com/v1",
    api_key=api_key
)

# import asyncio
#
# chain = (
#         llm | JsonOutputParser()
# )  # Due to a bug in older versions of Langchain, JsonOutputParser did not stream results from some models
#
#
# async def main():
#     async for text in chain.astream(
#             "output a list of the countries france, spain and japan and their populations in JSON format. "
#             'Use a dict with an outer key of "countries" which contains a list of countries. '
#             "Each country should have the key `name` and `population`"
#     ):
#         print(text, flush=True)
#
#
# asyncio.run(main())

# A function that operates on finalized inputs
# rather than on an input_stream
def _extract_country_names(inputs):
    """A function that does not operates on input streams and breaks streaming."""
    if not isinstance(inputs, dict):
        return ""

    if "countries" not in inputs:
        return ""

    countries = inputs["countries"]

    if not isinstance(countries, list):
        return ""

    country_names = [
        country.get("name") for country in countries if isinstance(country, dict)
    ]
    return country_names

chain = llm | JsonOutputParser() | _extract_country_names
async def main():
    async for text in chain.astream(
        "output a list of the countries france, spain and japan and their populations in JSON format. "
        'Use a dict with an outer key of "countries" which contains a list of countries. '
        "Each country should have the key `name` and `population`"
    ):
        print(text, end="|", flush=True)

asyncio.run(main())

